# -*- coding: utf-8 -*-

from deep_sort.deep_sort import nn_matching
from deep_sort.application_util import preprocessing
from deep_sort.deep_sort.detection import Detection
from deep_sort.deep_sort.tracker import Tracker
from deep_sort.tools import generate_detections
import numpy as np

class DeepSORT(object):
    _defaults = {
        "feature_model_filename": "resources/networks/mars-small128.pb",
        # "min_confidence": 0.8, # use YOLO score param
        # "nms_max_overlap": 1.0, # follow YOLO result
        # "min_detection_height": 0, # don't use this param, follow YOLO result
        "max_cosine_distance": 0.2,
        "nn_budget": None,
        "update_step": 1
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, detector, **kwargs):
        self.__dict__.update(self._defaults)  # set up default values
        self.__dict__.update(kwargs)  # and update with user overrides

        self.detector = detector

        self.encoder = generate_detections.create_box_encoder(self.feature_model_filename, batch_size=32)

        self.metric = nn_matching.NearestNeighborDistanceMetric(
            "cosine", self.max_cosine_distance, self.nn_budget)
        self.tracker = Tracker(self.metric, update_step=self.update_step)
        self.frame_index = 0
        # keep filtered detect result
        self.bboxes = []
        self.scores = []
        self.classes = []

    def process_frame(self, frame):
        '''
        :param frame: opencv mat image
        :param bboxes: bounding boxes of detection result (x,y,w,h)
        :return:
        '''

        need_update = ((self.frame_index % self.update_step) == 0)

        self.tracker.predict()

        if need_update:
            self.bboxes, self.scores, self.classes = self.detector.detect_image(frame)
            features = self.encoder(frame, self.bboxes)
            detections = [Detection(bbox, 1.0, feature) for bbox, feature in zip(self.bboxes, features)]
            # Run non-maxima suppression.
            # boxes = np.array([d.tlwh for d in detections])
            # scores = np.array([d.confidence for d in detections])
            # indices = preprocessing.non_max_suppression(boxes, self.nms_max_overlap, scores)
            # detections = [detections[i] for i in indices]
            self.tracker.update(detections)

        # Store results.
        results = []
        for track in self.tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > self.update_step:
                continue
            bbox = track.to_tlwh()
            results.append([
                self.frame_index, track.track_id, int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])])

        self.frame_index += 1

        return results
